#!/usr/bin/env python3
# coding=utf-8
# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
import getpass
import json
import subprocess
import time
import logging
import os
from ssl import create_default_context, Purpose
import urllib3
from enum import Enum
import argparse


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()  # 同时输出到控制台
    ]
)


class ClusterStatus(Enum):
    HEALTHY = 1
    ERROR = 2
    ACCURACY_ERROR = 3


def kubectl_get_pods_info():
    return subprocess.run(
        ["kubectl", "get", "pods", "-A", "-owide"],
        capture_output=True,
        text=True,
        check=True
    ).stdout


def load_cert():
    context = create_default_context(Purpose.SERVER_AUTH)
    cert_file_map = {
        "ca_cert": "./security/ca.pem",
        "tls_cert": "./security/cert.pem",
        "tls_key": "./security/cert.key.pem",
    }
    for _, cert_files in cert_file_map.items():
        if not os.path.exists(cert_files):
            return None
    password = getpass.getpass("Please enter the coordinator cert password: ")
    context.load_verify_locations(cafile=cert_file_map["ca_cert"])
    context.load_cert_chain(
        certfile=cert_file_map["tls_cert"],
        keyfile=cert_file_map["tls_key"],
        password=password
    )
    return context


def fetch_server_ip_and_port():
    pods_info = kubectl_get_pods_info()

    def fetch_ip_with_name(name: str):
        pod_info_lines = pods_info.split("\n")
        ip_idx = pod_info_lines[0].find("IP")
        for line in pod_info_lines:
            if name in line:
                return line[ip_idx:].split(" ")[0].strip()
        return ""

    def fetch_port(file: str):
        with open(file) as f:
            for line in f:
                if "port:" in line:
                    return line.split("port:")[-1].strip()
        return ""

    haproxy_ip = fetch_ip_with_name("haproxy")
    if haproxy_ip:
        return f"{haproxy_ip}:{fetch_port('./deployment/haproxy_init.yaml')}"
    return f"{fetch_ip_with_name('mindie-coordinator')}:{fetch_port('./deployment/mindie_ms_coordinator.yaml')}"


def check_service_status(
    http_pool_manager, 
    with_cert: bool, 
    enable_acc_check: bool,
    model_name: str, 
    input_content: str, 
    golden_result: str
):
    try:
        ip_and_port = fetch_server_ip_and_port()
        logging.info(f"Fetch server ip and port successfully: {ip_and_port}")
        http_prefix = "https" if with_cert else "http"
        response = http_pool_manager.request(
            "POST",
            f"{http_prefix}://{ip_and_port}/v1/completions",
            headers={"Content-Type": "application/json"},
            body=json.dumps({
                "model": model_name,
                "prompt": input_content,
                "temperature": 0,
                "max_new_tokens": 10,
                "stream": False,
            }).encode())
        if response.status >= 400:
            logging.info(f"Response from Coordinator failed, status is {response.status}, "
                         f"content is {response.data.decode()}")
            return ClusterStatus.ERROR
        resp_text = response.data.decode(errors="ignore")
        if enable_acc_check and golden_result not in resp_text:
            logging.info(f"Accuracy error: golden_result '{golden_result}' not in response: {resp_text}")
            return ClusterStatus.ACCURACY_ERROR
    except Exception as e:
        logging.info(f"Failed to connect to coordinator because {e}")
        return ClusterStatus.ERROR
    logging.info("MindIE MS Coordinator is ready!!!")
    return ClusterStatus.HEALTHY


def is_mindie_service_detected() -> bool:
    pod_status_info = kubectl_get_pods_info()
    return ("mindie-controller" in pod_status_info or "mindie-coordinator" in pod_status_info or
            "mindie-server" in pod_status_info)


def graceful_exit() -> None:
    subprocess.run(["bash", "show_logs.sh"])
    if not os.path.exists(os.path.join(os.getcwd(), "delete.sh")):
        raise RuntimeError("delete.sh not found, couldn't exit gracefully!!!")
    subprocess.run(["bash", "delete.sh"])
    while True:
        if not is_mindie_service_detected():
            logging.info("Delete mindie subprocess successfully!")
            return
        logging.info("Waiting for mindie subprocess to terminate!!!")
        time.sleep(10)


def start_service(boot_args) -> int:
    deploy_ac_job_res = subprocess.run(["python3", "deploy_ac_job.py"] + boot_args)
    if deploy_ac_job_res.returncode:
        logging.error(f"Start service failed! please check boot_args: {boot_args}")
        return 1
    else:
        logging.info(f"Start service successfully!")
        return 0


def main() -> None:
    parser = argparse.ArgumentParser(description="MindIE RAS Starter")
    parser.add_argument('--attach', action='store_true', help='是否是附加启动，如果是则不会删除原有服务')
    parser.add_argument('--enable-acc-check', action='store_true', help='是否启动精度检测')
    parser.add_argument('--model-name', type=str, default="dsv3_w8a8", help='模型名称')
    parser.add_argument('--probe-interval', type=int, default=10, help='探测周期（秒），默认10秒')
    parser.add_argument('--input-content', type=str, default='不要思考只回答他的名字，相对论的发明者是谁？', help='输入内容')
    parser.add_argument('--golden-result', type=str, default='爱因斯坦', help='标准结果')
    parser.add_argument('--max-fail-count', type=int, default=3, help='一个精度探测迭代内允许的最大失败次数，默认3次')
    parser.add_argument('--iter-dial-count', type=int, default=12, help='一个精度探测迭代的次数，默认12次')
    parser.add_argument('--max-unavailable-time', type=int, default=1200, help='最大服务不可用时长，默认1200秒')
    parser.add_argument('--http-timeout', type=int, default=60, help='urllib3请求超时时间，默认60秒')
    parser.add_argument('--restart', type=str, default='true', choices=['true', 'false'], help='是否重启，默认重启')
    args, boot_args = parser.parse_known_args()

    logging.info(f"启动参数: {vars(args)}, {boot_args}")

    enable_acc_check = args.enable_acc_check
    model_name = args.model_name
    probe_interval = args.probe_interval
    input_content = args.input_content
    golden_result = args.golden_result
    max_fail_count = args.max_fail_count
    iter_dial_count = args.iter_dial_count
    max_unavailable_time = args.max_unavailable_time
    http_timeout = args.http_timeout
    should_restart = args.restart.lower() == 'true'
    is_attach = args.attach

    if not is_attach and is_mindie_service_detected():
        logging.info("MindIE Service is already existed before start, restore all logs and delete it!")
        graceful_exit()

    cert_context = load_cert()
    if cert_context:
        logging.info("Sending requests to Coordinator with ssl!")
        http_pool_manager = urllib3.PoolManager(ssl_context=cert_context,
                                                assert_hostname=False,
                                                timeout=http_timeout,
                                                retries=False)
    else:
        logging.info("Sending requests to Coordinator without ssl!")
        http_pool_manager = urllib3.PoolManager(cert_reqs="CERT_NONE",
                                                timeout=http_timeout,
                                                retries=False)

    # Only start service in start up mode
    if not is_attach:
        result = start_service(boot_args)
        if result == 1: # just exit, need to check boot_args
            graceful_exit()
            return

    max_retry_time = 10240
    have_recorded_log = False
    while max_retry_time > 0:

        must_restart = False
        last_success_timepoint = time.time()
        while True:
            acc_fail_cnt, dial_cnt = 0, 0
            while dial_cnt < iter_dial_count:
                status = check_service_status(http_pool_manager, 
                                              cert_context is not None, 
                                              enable_acc_check,
                                              model_name,
                                              input_content, 
                                              golden_result)
                if status == ClusterStatus.HEALTHY:
                    last_success_timepoint = time.time()
                elif status == ClusterStatus.ERROR:
                    unavailable_time = time.time() - last_success_timepoint
                    logging.info(f"Unavailable time {unavailable_time}")
                    if unavailable_time > max_unavailable_time:
                        logging.error(f"Service unavailable time reach {max_unavailable_time}, need to restart!")
                        must_restart = True
                        break
                elif status == ClusterStatus.ACCURACY_ERROR:
                    acc_fail_cnt += 1
                    logging.info(f"Probe failed, status: {status}, fail_times: {acc_fail_cnt}")
                if acc_fail_cnt >= max_fail_count:
                    logging.info(f"Fail times {acc_fail_cnt} >= {max_fail_count}, will restart service!")
                    break
                dial_cnt += 1
                time.sleep(probe_interval)

            if must_restart: # 挂死检测退出, 不复用精度检测的不重启标志，触发了挂死超时必须重启
                break
            
            if acc_fail_cnt >= max_fail_count: # 精度检测退出
                if should_restart:
                    logging.error("Detected service accuracy is wrong, restore all "
                                  "logs and restart the whole service!")
                    must_restart = True
                    break
                else:
                    logging.error("Detected service accuracy is wrong, but not restarting "
                                  "as config, just record the logs.")
                    if not have_recorded_log:
                        subprocess.run(["bash", "show_logs.sh"])
                        have_recorded_log = True

        if must_restart:
            logging.info(f"Restarting service....")
            graceful_exit()

            result = start_service(boot_args)
            if result == 1:
                graceful_exit()
                return

        max_retry_time -= 1


if __name__ == '__main__':
    main()